package androidx.lifecycle

/**
 * @author : Shonn.Zhang
 * @email  : shonn.zhang@mart.blue
 * @time   : 2021/6/25 11:42
 * @desc   :
 */
class ExtendLiveData<T> : MutableLiveData<T>() {

    private val mExtendObservers = mutableMapOf<Observer<in T>, ExtendObserver<T>>()

    override fun observe(owner: LifecycleOwner, observer: Observer<in T>) {
        val extendObserver = mExtendObservers.getOrPut(observer) {
            ExtendObserver<T>(this, observer).apply {
                mExtendObservers[observer] = this
            }
        }
        super.observe(owner, extendObserver)
    }

    override fun observeForever(observer: Observer<in T>) {
        val extendObserver = mExtendObservers.getOrPut(observer) {
            ExtendObserver(this, observer).apply {
                mExtendObservers[observer] = this
            }
        }
        super.observeForever(extendObserver)
    }

    override fun removeObserver(observer: Observer<in T>) {
        val removed = mExtendObservers.remove(observer) ?: observer.apply {
            if (observer is ExtendObserver) {
                mExtendObservers.remove(observer.observer)
            }
        }
        super.removeObserver(removed)
    }

}